#include "compi/compi.h"

#include <iostream>
#include <mpi.h>
#include <string>
#include <vector>

using namespace compi;

// 演示：不协商标签，直接使用命名空间过滤（ANY_TAG）的 int 点对点收发。
// 发送侧使用 SendRequest（自动管理命名空间标签）；接收侧使用 RecvRequest（ANY_TAG 过滤）。
static void
run_ping_pong_int_any(Context & comm, int rank, TagNamespace ns)
{
    if (rank == 0)
    {
        int value = 123;
        int reply = 0;

        // 构造请求：发送到 rank 1；接收来自 rank 1（ANY_TAG + 命名空间过滤）
        SendRequest sreq(comm, /*dest=*/1, &value, MPI_INT, ns);
        RecvRequest rreq(comm, /*src=*/1, &reply, MPI_INT, ns);

        std::cout << "rank 0 posted send int: " << value << std::endl;

        // 使用进度引擎，设置最大在途窗口为 2
        ProgressEngine engine(/*max_inflight=*/2);
        engine.add(std::move(sreq));
        engine.add(std::move(rreq));

        // 等待两个请求都完成
        while (engine.waitsome() < 2)
            ;

        std::cout << "rank 0 received int: " << reply << std::endl;
    }
    else if (rank == 1)
    {
        int received = 0;

        // 先接收来自 rank 0 的消息
        RecvRequest rreq(comm, /*src=*/0, &received, MPI_INT, ns);
        ProgressEngine engine_recv(/*max_inflight=*/1);
        engine_recv.add(std::move(rreq));
        while (engine_recv.waitsome() < 1)
            ;
        std::cout << "rank 1 received int: " << received << std::endl;

        // 计算响应并发送回 rank 0
        int response = received + 1;
        SendRequest sreq(comm, /*dest=*/0, &response, MPI_INT, ns);
        ProgressEngine engine_send(/*max_inflight=*/1);
        engine_send.add(std::move(sreq));
        while (engine_send.waitsome() < 1)
            ;
        std::cout << "rank 1 sent int: " << response << std::endl;
    }
}

// 演示：字符串的 ANY_TAG + 命名空间过滤的非固定长度收发（序列化路径）。
static void
run_ping_pong_string_any(Context & comm, int rank, TagNamespace ns)
{
    if (rank == 0)
    {
        std::string msg = "hello";
        std::string reply;

        SendRequest sreq(comm, /*dest=*/1, msg, ns);
        RecvRequest rreq(comm, /*src=*/1, reply, ns);

        std::cout << "rank 0 posted send string: " << msg << std::endl;

        ProgressEngine engine(/*max_inflight=*/2);
        engine.add(std::move(sreq));
        engine.add(std::move(rreq));
        while (engine.waitsome() < 2)
            ;

        std::cout << "rank 0 received string: " << reply << std::endl;
    }
    else if (rank == 1)
    {
        std::string received;
        RecvRequest rreq(comm, /*src=*/0, received, ns);
        ProgressEngine engine_recv(/*max_inflight=*/1);
        engine_recv.add(std::move(rreq));
        while (engine_recv.waitsome() < 1)
            ;
        std::cout << "rank 1 received string: " << received << std::endl;

        std::string response = received + " world";
        SendRequest sreq(comm, /*dest=*/0, response, ns);
        ProgressEngine engine_send(/*max_inflight=*/1);
        engine_send.add(std::move(sreq));
        while (engine_send.waitsome() < 1)
            ;
        std::cout << "rank 1 sent string: " << response << std::endl;
    }
}

// 演示：字节向量的 ANY_TAG + 命名空间过滤非固定长度收发（序列化路径）。
static void
run_ping_pong_bytes_any(Context & comm, int rank, TagNamespace ns)
{
    auto to_bytes = [](const std::string & s)
    {
        std::vector<std::byte> v(s.size());
        std::memcpy(v.data(), s.data(), s.size());
        return v;
    };
    auto to_string = [](const std::vector<std::byte> & v)
    { return std::string(reinterpret_cast<const char *>(v.data()), v.size()); };

    if (rank == 0)
    {
        std::vector<std::byte> payload = to_bytes("ping");
        std::vector<std::byte> reply;

        SendRequest sreq(comm, /*dest=*/1, payload, ns);
        RecvRequest rreq(comm, /*src=*/1, reply, ns);

        std::cout << "rank 0 posted send bytes: " << to_string(payload) << " (" << payload.size()
                  << " bytes)" << std::endl;

        ProgressEngine engine(/*max_inflight=*/2);
        engine.add(std::move(sreq));
        engine.add(std::move(rreq));
        while (engine.waitsome() < 2)
            ;

        std::cout << "rank 0 received bytes: " << to_string(reply) << " (" << reply.size()
                  << " bytes)" << std::endl;
    }
    else if (rank == 1)
    {
        std::vector<std::byte> received;
        RecvRequest rreq(comm, /*src=*/0, received, ns);
        ProgressEngine engine_recv(/*max_inflight=*/1);
        engine_recv.add(std::move(rreq));
        while (engine_recv.waitsome() < 1)
            ;
        std::cout << "rank 1 received bytes: " << to_string(received) << " (" << received.size()
                  << " bytes)" << std::endl;

        // 生成响应：在尾部追加 "!"
        std::vector<std::byte> response = received;
        response.push_back(static_cast<std::byte>('!'));
        SendRequest sreq(comm, /*dest=*/0, response, ns);
        ProgressEngine engine_send(/*max_inflight=*/1);
        engine_send.add(std::move(sreq));
        while (engine_send.waitsome() < 1)
            ;
        std::cout << "rank 1 sent bytes: " << to_string(response) << " (" << response.size()
                  << " bytes)" << std::endl;
    }
}

// 程序入口：要求使用 mpirun/mpiexec 以至少 2 个进程运行。
int
main()
{
    Environment env;
    env.set_failure_policy(FailurePolicy::ThrowException);
    Context & comm = Context::for_comm(MPI_COMM_WORLD, env);

    const int rank = comm.rank();
    const int world_size = comm.size();
    if (world_size < 2)
    {
        if (rank == 0)
        {
            std::cerr << "This example requires at least 2 MPI processes." << std::endl;
        }
        return 0;
    }

    // 创建示例命名空间，用于隔离示例中的消息标签
    auto & alloc = comm.tag_allocator();
    auto ns = alloc.create_namespace("example");

    // 使用新接口的 ANY_TAG 过滤收发（固定类型 int）
    run_ping_pong_int_any(comm, rank, ns);

    // 字符串序列化路径示例
    run_ping_pong_string_any(comm, rank, ns);

    // 字节向量序列化路径示例
    run_ping_pong_bytes_any(comm, rank, ns);

    return 0;
}
